#include <ctrl_alg/lqr/lqr.hpp>
#include <general_interface/msg/control_signal.hpp>
#include <rclcpp/rclcpp.hpp>
#include <sensor_msgs/msg/imu.hpp>
#include <std_msgs/msg/float64.hpp>
#include <tf2/LinearMath/Matrix3x3.h>
#include <tf2/LinearMath/Quaternion.h>

static const int control_rate = 50;
static const int state_publish_rate = 50;

class ControlNode : public rclcpp::Node
{
  private:
    using WheelSpeed = std_msgs::msg::Float64;
    rclcpp::Subscription<WheelSpeed>::SharedPtr leftSpeedSub, rightSpeedSub;
    using WheelTorque = std_msgs::msg::Float64;
    rclcpp::Publisher<WheelTorque>::SharedPtr leftTorquePub, rightTorquePub;
    rclcpp::Subscription<sensor_msgs::msg::Imu>::SharedPtr imuSub;
    rclcpp::Subscription<general_interface::msg::ControlSignal>::SharedPtr
        controlSub;

    std::unique_ptr<CtrlAlg::LQR<double, 6, 2>> lqr;
    rclcpp::TimerBase::SharedPtr timer_balance;
    rclcpp::TimerBase::SharedPtr timer_publish_state;
    OnSetParametersCallbackHandle::SharedPtr param_callback_handle;

    std::map<std::string, rclcpp::Publisher<std_msgs::msg::Float64>::SharedPtr>
        publishers;

    struct State
    {
        double gyroy = 0;
        double gyroz = 0;
        double yaw = 0;
        double pitch = 0;
        double roll = 0;
        double speed = 0;

        double leftCurrent = 0;
        double rightCurrent = 0;

        double leftRPM = 0;
        double rightRPM = 0;

        double targetSpeed = 0;
        double targetAngle = 0;
        double targetRotate = 0;
    };
    State state;

  public:
    ControlNode() : Node("control_node")
    {
        RCLCPP_INFO(get_logger(), "load parameters for PID");
        param_callback_handle = RegeistParamCallBack();
        RCLCPP_INFO(get_logger(), "PID parameters loaded");

        RCLCPP_INFO(get_logger(), "Creat ROS2 related stuff");
        leftSpeedSub = create_subscription<WheelSpeed>(
            "left_speed", 1,
            [this](WheelSpeed::SharedPtr msg) { state.leftRPM = msg->data; });
        rightSpeedSub = create_subscription<WheelSpeed>(
            "right_speed", 1,
            [this](WheelSpeed::SharedPtr msg) { state.rightRPM = msg->data; });
        leftTorquePub = create_publisher<WheelTorque>("left_torque", 1);
        rightTorquePub = create_publisher<WheelTorque>("right_torque", 1);
        imuSub = create_subscription<sensor_msgs::msg::Imu>(
            "imu", 1,
            [this](sensor_msgs::msg::Imu::SharedPtr msg)
            {
                state.gyroy =
                    -msg->angular_velocity.y * std::numbers::inv_pi * 180;
                state.gyroz =
                    msg->angular_velocity.z * std::numbers::inv_pi * 180;
                tf2::Quaternion q(msg->orientation.x, msg->orientation.y,
                                  msg->orientation.z, msg->orientation.w);
                tf2::Matrix3x3 m(q);
                m.getRPY(state.roll, state.pitch, state.yaw);
                state.roll *= std::numbers::inv_pi * 180;
                state.pitch *= std::numbers::inv_pi * 180;
                state.yaw *= std::numbers::inv_pi * 180;
            });
        controlSub = create_subscription<general_interface::msg::ControlSignal>(
            "control_signal", 1,
            [this](general_interface::msg::ControlSignal::SharedPtr msg)
            {
                state.targetRotate =
                    std::clamp(msg->rotate_control * std::numbers::inv_pi * 180,
                               -10.0, 10.0);
                state.targetAngle = std::clamp(msg->forward_control *
                                                   std::numbers::inv_pi * 180,
                                               -20.0, 20.0);
            });
        RCLCPP_INFO(get_logger(), "Finish creating ROS2 related stuff");

        RCLCPP_INFO(get_logger(), "creat timer for control loop");
        timer_balance =
            create_wall_timer(std::chrono::milliseconds(1000 / control_rate),
                              std::bind(&ControlNode::OnLoop, this));
        timer_publish_state = create_wall_timer(
            std::chrono::milliseconds(1000 / state_publish_rate),
            std::bind(&ControlNode::OnPublishState, this));
        RCLCPP_INFO(get_logger(), "timer for control loop created");
    }

  private:
    OnSetParametersCallbackHandle::SharedPtr RegeistParamCallBack()
    {
        return add_on_set_parameters_callback(
            [&](const std::vector<rclcpp::Parameter> &parameters)
            {
                rcl_interfaces::msg::SetParametersResult result;
                result.successful = true;
                for (const auto &param : parameters)
                {
                    if (param.get_name() == "gyro_pid")
                    {
                        
                    }
                    else
                    {
                        result.successful = false;
                        result.reason =
                            "unknown parameter name: " + param.get_name();
                        break;
                    }
                }
                return result;
            });
    }

    void OnLoop()
    {
        //state.leftCurrent =
        //    ClampAbs(state.forwardTorque - state.rotateTorque, 80);
        //state.rightCurrent =
        //    ClampAbs(state.forwardTorque + state.rotateTorque, 80);

        state.speed = (state.leftRPM + state.rightRPM) / 2;

        WheelSpeed msg;
        msg.data = state.leftCurrent;
        leftTorquePub->publish(msg);
        msg.data = state.rightCurrent;
        rightTorquePub->publish(msg);
    }

    void OnPublishState()
    {
        Publish("target_angle", state.targetAngle);
        Publish("gyro", state.gyroy);
        Publish("angle", state.pitch);
        Publish("target_gyroz", state.targetRotate);
        Publish("gyroz", state.gyroz);
        Publish("target_speed", state.targetSpeed);
        Publish("speed", state.speed);
    }

    void Publish(std::string topic, double value)
    {
        if (publishers.find(topic) == publishers.end())
            publishers[topic] =
                create_publisher<std_msgs::msg::Float64>(topic, 1);
        auto msg = std_msgs::msg::Float64();
        msg.data = value;
        publishers[topic]->publish(msg);
    }

    inline double ClampAbs(double value, double max)
    {
        return std::clamp(value, -max, max);
    }

    bool InRange(double value, double min, double max)
    {
        return value >= min && value <= max;
    }
};

int main(int argc, char **argv)
{
    rclcpp::init(argc, argv);
    rclcpp::spin(std::make_shared<ControlNode>());
    rclcpp::shutdown();
    return 0;
}